import matplotlib.pyplot as plt
import pickle

# 定义文本框和箭头格式。两种格式的文本框分别代表叶结点和其他节点。
decisionNode = dict(boxstyle="sawtooth", fc="0.8")  # 决策树根结点和内部结点的属性。
# boxstyle="sawtooth" ：文本框的边缘是波浪线，fc="0.8" ：颜色深度。可以写为decisionNode={boxstyle:'sawtooth',fc:'0.8'}
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")  # 箭头样式
# 下面两句解决中文乱码问题
plt.rcParams['font.sans-serif'] = [u'SimHei']
plt.rcParams['axes.unicode_minus'] = False


# 计算树的叶子节点数量，以确定x轴的长度
def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = list(myTree.keys())[0]  # dict.keys()返回的不在是list类型了，也不支持索引。先转换成list
    # firstStr = myTree.keys()[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':  # {'纹理': {'模糊': '坏瓜', '稍糊': {'触感': {'软粘': '好瓜',}}}}
            numLeafs += getNumLeafs(secondDict[key])  # 只要value字典，就一直递归，直到不是(模糊: 坏瓜),此时是叶结点+1
        else:
            numLeafs += 1
    return numLeafs


# 计算树的最大深度，以确定y轴的长度
def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[key]).__name__ == 'dict':  # 如果value是字典就加以，直到不是，取最大
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:
            thisDepth = 1
        if thisDepth > maxDepth:
            maxDepth = thisDepth
    return maxDepth


# 绘制带箭头的注解
def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction', \
                            xytext=centerPt, textcoords='axes fraction', va="center", ha="center", \
                            bbox=nodeType, arrowprops=arrow_args)


# nodeTxt为要显示的文本，centerPt为文本的中心点，parentPt为箭头指向文本的点，xy是箭头尖的坐标，xytest设置注释内容显示的中心位置
# xycoords和textcoords是坐标xy与xytext的说明（按轴坐标），若textcoords=None，则默认textcoords与xycoords相同，若都未设置，默认为data
# va/ha设置节点框中文字的位置，va为纵向取值为(u'top', u'bottom', u'center', u'baseline')，ha为横向取值为(u'center', u'right', u'left')


# 在父子节点间 填充文本信息。也就是剪头线上的标注，在两个点坐标的中心处添加标注
def plotMidText(cntrPt, parentPt, txtString):
    xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
    yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
    # print(xMid, yMid, txtString)
    createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


# 递归函数，树的宽度来计算判断节点的位置，是放在所有叶子节点的中间，不仅仅是它子结点的中间。
# 其中两个全局变量plotTree.xOff和plotTree.yOff，用于追踪已绘制的节点位置，并放置下个节点的恰当位置
def plotTree(myTree, parentPt, nodeTxt):  # nodeTxt:key
    numLeafs = getNumLeafs(myTree)
    depth = getTreeDepth(myTree)
    firstStr = list(myTree.keys())[0]
    cntrPt = (plotTree.xOff + (1.0 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yOff)  # 按照叶子结点个数划分x轴(0.5, 1.0)
    # print('cntrPt:', cntrPt)
    plotMidText(cntrPt, parentPt, nodeTxt)  # 箭头线上文本
    plotNode(firstStr, cntrPt, parentPt, decisionNode)  # 判断结点
    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD  # 采用的自顶向下的绘图，需要依次递减Y坐标；按比例减少plotTree.yOff
    for key in secondDict.keys():  # value为字典，继续递归
        if type(secondDict[key]).__name__ == 'dict':
            plotTree(secondDict[key], cntrPt, str(key))
        else:  # 为叶子结点
            plotTree.xOff = plotTree.xOff + 1.0 / plotTree.totalW  # x方向计算节点坐标
            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)  # 绘图，叶子结点
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))  # 添加文本信息
    plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD  # 下次重新调用时恢复y


# 主函数
def createPlot(inTree):
    fig = plt.figure(1, facecolor='white')  # 新建一个画布，背景设置为白色的
    fig.clf()  # 清空绘图区
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)  # frameon表示是否绘制坐标轴矩形
    plotTree.totalW = float(getNumLeafs(inTree))  # 树的宽度#全局变量宽度 = 叶子数
    plotTree.totalD = float(getTreeDepth(inTree))  # 树的深度#全局变量高度 = 深度
    plotTree.xOff = -0.5 / plotTree.totalW  # 例如绘制3个叶子结点，坐标应为1/3,2/3,3/3
    plotTree.yOff = 1.0  # 但这样会使整个图形偏右因此初始的，将x值向左移一点。
    plotTree(inTree, (0.5, 1.0), '')
    plt.show()


if __name__ == '__main__':
    filename = './DTModel/ID3Model'
    fr = open(filename, 'rb')
    tree = pickle.load(fr)
    createPlot(tree)
